import os
from urllib.request import urlretrieve

import numpy as np
import uvicorn
from fastapi import FastAPI, Form, File, UploadFile
from milvus import Milvus
from numpy import linalg
from tensorflow.keras.applications.resnet50 import ResNet50
from tensorflow.keras.applications.resnet50 import preprocess_input
from tensorflow.keras.preprocessing import image

# 环境变量
MILVUS_HOST = os.getenv("MILVUS_HOST", "192.168.1.10")
MILVUS_PORT = os.getenv("MILVUS_PORT", 19530)
MILVUS_NAME = os.getenv("MILVUS_NAME", "image_search")
UPLOAD_PATH = os.getenv("UPLOAD_PATH", "images")

# 创建目录
if not os.path.exists(UPLOAD_PATH):
    os.makedirs(UPLOAD_PATH)

# 创建对象
app = FastAPI()
model = ResNet50(weights="imagenet", pooling="max", include_top=False)
client = Milvus(MILVUS_HOST, MILVUS_PORT)


# 提取特征
def extract_features(img_path):
    img = image.load_img(img_path, target_size=(224, 224))
    img = image.img_to_array(img)
    img = np.expand_dims(img, axis=0)
    img = preprocess_input(img)
    feat = model.predict(img)
    norm_feat = feat[0] / linalg.norm(feat[0])
    return norm_feat.tolist()


# 插入向量
def insert_vectors(vectors):
    status, ids = client.insert(MILVUS_NAME, vectors)
    if status.code:
        raise Exception(status.message)
    return ids


# 搜索向量
def search_vectors(vectors, top_k):
    status, result = client.search(MILVUS_NAME, top_k, vectors, params={'nprobe': 16})
    if status.code:
        raise Exception(status.message)
    return result


# 上传图片
@app.post("/upload")
def upload_images(url: str = Form(...)):
    try:
        img_path = os.path.join(UPLOAD_PATH, os.path.basename(url))
        urlretrieve(url, img_path)
        features = extract_features(img_path)
        os.remove(img_path)
        ids = insert_vectors([features])
        return ids[0]
    except Exception as e:
        return {'status': False, 'msg': str(e)}, 400


# 搜索图片
@app.post("/search")
async def search_images(file: UploadFile = File(...), top_k: int = Form(50)):
    try:
        content = await file.read()
        img_path = os.path.join(UPLOAD_PATH, file.filename)
        with open(img_path, "wb+") as f:
            f.write(content)
        features = extract_features(img_path)
        os.remove(img_path)
        vectors = search_vectors([features], top_k)
        return [x.id for x in vectors[0]]
    except Exception as e:
        return {'status': False, 'msg': str(e)}, 400


# 启动服务
if __name__ == '__main__':
    uvicorn.run(app=app, host='0.0.0.0', port=5000)